# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

import collections
import enum
import functools
import json
import math
import typing
import uuid

import requests
import tornado.gen

import coin.flow.topic_map as topic_map
import coin.flow.subscriber as flow_subscriber
import coin.proto.coin_market_enums_pb2 as coin_enum
import coin.proto.coin_order_enums_pb2 as coin_order

from coin.base.config import Config
from coin.base.timestamp import get_timestamp
from coin.exchange.base.order_fill_util import FillManager
from coin.exchange.base.order_gateway_ready import SpotOrderGatewayReady
from coin.exchange.base.rest.task_scheduler import Scheduler, wrap_async_task
from coin.exchange.base.rest_order_gateway_base import PrivateExchangeMessage
from coin.exchange.gdax.kr_rest.private_client import GdaxPrivateParser
from coin.exchange.gdax.kr_rest.product import GdaxProduct, GdaxCurrency
from coin.exchange.util.rest_rate_limit import RestRateLimiter
from coin.proto.coin_order_gateway_pb2 import OrderGatewayConfig, OrderEvent
from coin.proto.coin_query_pb2 import AccountBalance
from coin.strategy.hk_hard_arb.og_util import print_working_orders

from coin.exchange.gdax.ws.ws_private_client import (
    GdaxWsPrivateClient,
    ChannelType,
    MessageType,
)
from coin.exchange.gdax.kr_rest.native_private_client import (
    GdaxOrderStatus,
    GdaxOrderSide,
    GdaxTimeInForce,
    GdaxNativePrivateClient,
    GdaxOrderType,
    GdaxStpFlag,
    GdaxFillType,
)
from coin.exchange.base.order_gateway import (
    BaseOrderGatewayInfoMutator,
    Order,
    OrderGatewayStatus,
    OrderType,
    OrderSide,
)
from coin.exchange.base.rest_order_gateway_base import (
    RestOrderGateway,
    create_order_internal,
    exchange_message_handler,
)


def to_gdax_order_status(status, reason=None):
  status = GdaxOrderStatus(status)
  if status == GdaxOrderStatus.DONE:
    if reason == 'canceled':
      status = GdaxOrderStatus.DONE_CANCELED
    elif reason == 'filled':
      status = GdaxOrderStatus.DONE_FILLED
    else:
      raise ValueError('status=%s, reason=%s' % (status, reason))
  return status


def create_order_from_open_order_message(message, timestamp, order_id=None):
  """
  {
  "id": "5915a03b-090b-4275-be23-b41cf38134a0",
  "price": "134.00000000",
  "size": "0.100$0000",
  "product_id": "ETH-USD",
  "side": "sell",
  "type": "limit",
  "time_in_force": "GTC",
  "post_only": true,
  "created_at": "2019-03-14T07:30:30.314884Z",
  "fill_fees": "0.000000$000000000",
  "filled_size": "0.00000000",
  "executed_value": "0.0000000000000000",
  "status": "pending",
  "settled": false
  }
  """
  price_orig = float(message['price'])
  qty_orig = float(message['size'])
  product = GdaxProduct.FromStrNativeProduct(message['product_id'])
  exchange_order_id = message['id']
  order_side = getattr(OrderSide, GdaxOrderSide(message['side']).name)
  order_type = getattr(OrderType, GdaxOrderType(message['type']).name)
  qty_accum_fill = float(message['filled_size'])
  price_avg_fill = None
  if qty_accum_fill > 0:
    price_avg_fill = float(message['executed_value']) / qty_accum_fill
  fee = float(message['fill_fees'])
  status = to_gdax_order_status(message['status'], message.get('done_reason'))
  fill = FillManager(
      price_orig=price_orig,
      qty_orig=qty_orig,
      price_last_fill=price_avg_fill,
      qty_last_fill=qty_accum_fill,
      price_avg_fill=price_avg_fill,
      qty_accum_fill=qty_accum_fill,
      timestamp_last_fill=timestamp,
  )

  order = Order(order_id=order_id,
                product=product,
                price=price_orig,
                qty=qty_orig,
                order_side=order_side,
                order_type=order_type,
                tag=None,
                accepted=True,
                cancel_sent=False,
                internal=create_order_internal(
                    timestamp=timestamp,
                    exchange_order_id=exchange_order_id,
                    fill=fill,
                    fee=fee,
                    status=status,
                ))
  return order


def create_order_from_ws_order_update(message, timestamp, order_id=None):
  """
  {
    "type": "done",
    "side": "sell",
    "order_id": "uuid",
    "reason": "canceled",
    "product_id": "BTC-USD",
    "price": "26.00000000",
    "remaining_size": "0.02000000",
    "sequence": 42673252,
    "user_id": "5ab34e17b223870441a66250",
    "profile_id": "cc5671bf-860d-4fbd-9f23-a791fc6eeb88",
    "time": "2019-03-14T03:41:19.181000Z"
  }
  """
  price_orig = float(message['price'])
  qty_orig = float(message.get('size', 0)) or float(message.get('remaining_size', 0))
  product = GdaxProduct.FromStrNativeProduct(message['product_id'])
  exchange_order_id = message['order_id']
  order_side = getattr(OrderSide, GdaxOrderSide(message['side']).name)
  order_type = OrderType.LIMIT
  status = to_gdax_order_status(message['type'], message.get('reason'))
  fill = FillManager(
      price_orig=price_orig,
      qty_orig=qty_orig,
      timestamp_last_fill=timestamp,
  )
  order = Order(order_id=order_id,
                product=product,
                price=price_orig,
                qty=qty_orig,
                order_side=order_side,
                order_type=order_type,
                tag=None,
                accepted=True,
                cancel_sent=False,
                internal=create_order_internal(
                    timestamp=timestamp,
                    exchange_order_id=exchange_order_id,
                    fill=fill,
                    status=status,
                ))
  return order


class GdaxOrderGatewayInfo(BaseOrderGatewayInfoMutator):
  CurrencyType = GdaxCurrency
  ProductType = GdaxProduct
  BalanceType = float
  PositionType = float

  def __init__(self, currencies: typing.List[GdaxCurrency], products: typing.List[GdaxProduct]):
    super().__init__(currencies, products)

  def update_balance(self, message):
    msg = GdaxPrivateParser.parse_account_balance(message)
    for balance in msg.each_balance:
      currency = GdaxCurrency.FromStrNativeCurrency(balance.currency_native)
      if currency in self.currencies:
        self.set_balance(currency, balance.total, balance)

  def balance_as_proto(self):
    return AccountBalance(exchange='Gdax', each_balance=self.get_notNone_balances_internal())


class GdaxOrderGateway(RestOrderGateway):
  _exchange = 'Gdax'
  _market_type = 'Spot'
  _api_version = None

  OGInfoType = GdaxOrderGatewayInfo

  class _MsgType(enum.IntEnum):
    REST_ACCOUNT_BALANCE = 1
    REST_ORDERS = 2
    REST_FILLS = 3
    WS_HEART_BEAT = 4
    WS_USER = 5

  def __init__(self,
               currencies: typing.List[GdaxCurrency],
               products: typing.List[GdaxProduct],
               config: Config,
               og_config=None,
               logger=None):
    super().__init__(currencies, products, config, og_config, logger)
    self._og_ready = SpotOrderGatewayReady()
    self._og_ready.set_status(OrderGatewayStatus.INITIALIZING)

    # Following instance attributes are initialized in start method.
    self._flow_sub = None
    self._ioloop = None
    self._publisher = None
    self._scheduler = None
    self._private_client = None
    self._ws_private_client = None
    self._rest_rate_limiter = RestRateLimiter(delay=10)
    self._processed_fills = set()  # For rest fills only

    self._og_config = OrderGatewayConfig(
        account_query_period=5,
        orders_snapshot_query_period=1,
        fill_query_period=3,
    )
    if og_config is not None:
      self._og_config.MergeFrom(og_config)
    self._og_logger.exchange_type = coin_enum.Gdax
    self._rest_fills = collections.defaultdict(FillManager)
    self._ws_fills = collections.defaultdict(FillManager)

  def submittable(self):
    return self.is_ready() and self._rest_rate_limiter.is_available()

  def start(self, queue, ioloop, flow_sub=None):
    api_host = self._og_config.rest_root_url if \
               self._og_config.HasField('rest_root_url') else None
    disable_ssl = self._og_config.rest_disable_ssl

    topic_id = topic_map.search_topic_id('gdax_rest_private_1')[0][1]
    self._ioloop = ioloop

    self._publisher = queue.get_writer(topic_id=topic_id)
    self._scheduler = Scheduler(ioloop=ioloop, logger=self._logger)

    self._private_client = GdaxNativePrivateClient(key_file=self._config.key_file,
                                                   api_host=api_host,
                                                   disable_ssl=disable_ssl,
                                                   use_async=True)
    self._prepare_ws_private_client(start=True)
    self._install_periodic_query()

    # Subscribe to flow for feed.
    self._flow_sub = flow_sub or flow_subscriber.from_queue(queue)
    self._flow_sub.subscribe('gdax_rest_private_1', self.on_private_exchange_msg)

  @tornado.gen.coroutine
  def _install_periodic_query(self):
    assert self._scheduler is not None, 'self._scheduler is not set!'
    assert self._publisher is not None, 'self._publisher is not set!'

    query_balance_task = wrap_async_task(
        self._private_client.query_accounts,
        msg_type=self._MsgType.REST_ACCOUNT_BALANCE,
        publisher=self._publisher,
    )
    self._scheduler.add_periodic_task(
        period=self._og_config.account_query_period,
        job_func=query_balance_task,
    )

    query_orders_task = wrap_async_task(self._private_client.query_orders,
                                        msg_type=self._MsgType.REST_ORDERS,
                                        publisher=self._publisher,
                                        func_params=dict(status=GdaxOrderStatus.ALL))
    self._scheduler.add_periodic_task(period=self._og_config.orders_snapshot_query_period,
                                      job_func=query_orders_task)

    for product in self.og_info.products:
      query_fills_task = wrap_async_task(self._private_client.query_fills,
                                         publisher=self._publisher,
                                         msg_type=self._MsgType.REST_FILLS,
                                         func_params=dict(product_id=product.native_symbol))
      self._scheduler.add_periodic_task(period=self._og_config.fill_query_period,
                                        job_func=query_fills_task)

    def print_info():
      print_working_orders(self._logger, self)

    self._scheduler.add_periodic_task(
        period=10,
        job_func=print_info,
    )
    self._scheduler.start()

  def _prepare_ws_private_client(self, start=True):
    if self._ws_private_client is not None:
      self._ws_private_client.stop()

    close_callback = functools.partial(self._prepare_ws_private_client, True)
    self._ws_private_client = GdaxWsPrivateClient(
        self._config.key_file,
        ioloop=self._ioloop,
        notification_callback=self._notification_callback,
        close_callback=close_callback,
        on_ready=self._on_ws_private_client_ready,
    )
    if start:
      self._ws_private_client.start()

  def _notification_callback(self, ws_msg):
    timestamp = get_timestamp()
    user_msg = (
        MessageType.RECEIVED,
        MessageType.OPEN,
        MessageType.CHANGE,
        MessageType.DONE,
        MessageType.MATCH,
    )
    ws_msg_type = MessageType(ws_msg['type'])
    if ws_msg_type == MessageType.HEARTBEAT:
      msg_type = self._MsgType.WS_HEART_BEAT
    elif ws_msg_type in user_msg:
      msg_type = self._MsgType.WS_USER
    else:
      self._logger.error('Unhandled message type: %s', ws_msg)
      return

    exchange_message = PrivateExchangeMessage(msg_type=msg_type,
                                              data=ws_msg,
                                              received_timestamp=timestamp,
                                              request_params=None,
                                              request_timestamp=None)
    self._publisher.write(timestamp, exchange_message)

  def _on_ws_private_client_ready(self):
    native_symbols = [prod.native_symbol for prod in self._og_info.products]
    channels = [ChannelType.USER, ChannelType.HEARTBEAT]
    self._ws_private_client.subscribe_auth_channels(native_symbols, channels)

  def _handle_error_response(self, response):
    if response.status_code in (502, 504):
      self._logger.error('Cloudflare bad gateway error! %s, %s!',
                         response.status_code,
                         response.url)
    elif response.status_code == 429:
      self._rest_rate_limiter.set_error(response.status_code)
      self._logger.error('Rate limit hit! %s, %s', response.status_code, response.url)
    else:
      self._logger.error('response error! %s, %s, %s, %s',
                         response.status_code,
                         response.content.decode(),
                         response.request.method,
                         response.url)

  def _async_submit_order(self, order, post_only):
    assert order.order_id is not None, order
    assert order.order_type == OrderType.LIMIT, order
    formatter = self._order_formatter[order.product]
    client_oid = str(uuid.uuid1())[:-1] + 'a'
    order.tag = client_oid
    params = {
        'product_id': order.product.native_symbol,
        'client_oid': client_oid,
        'side': GdaxOrderSide(order.order_side.name.lower()),
        'price': formatter.format_price(order.price),
        'size': formatter.format_qty(order.qty),
        'time_in_force': GdaxTimeInForce.GTC,
        'order_type': GdaxOrderType(order.order_type.name.lower()),
        'stp': GdaxStpFlag.DECREASE_AND_CANCEL,
        'post_only': post_only,
    }
    fut = self._private_client.place_order(**params)
    return fut

  def _async_cancel_order(self, order):
    assert order.internal.exchange_order_id is not None, order
    fut = self._private_client.cancel_order(order_id=order.internal.exchange_order_id)
    return fut

  def _handle_submit_response(self, order, response):
    """
    {
    "id": "e4e2bcc3-3489-464c-bf06-c287c50b105e",
    "price": "455.00000000",
    "size": "0.02000000",
    "product_id": "BTC-USD",
    "side": "buy",
    "stp": "dc",
    "type": "limit",
    "time_in_force": "GTC",
    "post_only": true,
    "created_at": "2019-03-08T09:16:45.051065Z",
    "fill_fees": "0.0000000000000000",
    "filled_size": "0.00000000",
    "executed_value": "0.0000000000000000",
    "status": "pending",
    "settled": false
    }
    """
    if response.status_code != requests.codes.ok:
      self._order_manager.remove_order(order, ignore_error=True)
      self._handle_error_response(response)
      return OrderEvent.ORDER_ERROR

    message = response.json()
    if message['status'] == 'rejected':
      self._logger.error('Order rejected %s', message)
      self._order_manager.remove_order(order, ignore_error=True)
      return OrderEvent.ORDER_REJECTED

    exchange_order_id = str(message['id'])
    prev_order = self._order_manager.get(exchange_order_id=exchange_order_id)
    if prev_order is None:
      self._order_manager.set_order(order, exchange_order_id=exchange_order_id)
    else:
      if prev_order.order_id != order.order_id:
        self._logger.info('order id replaced: %s => %s', prev_order.order_id, order.order_id)
        self._order_manager.remove_order(prev_order)
        self._order_manager.set_order(order, exchange_order_id=exchange_order_id)
        self._order_manager.remove_order(order)
        prev_order.order_id = order.order_id
        self._order_manager.add_order(prev_order)
    return OrderEvent.ORDER_ACCEPTED

  def _handle_cancel_response(self, order, response):
    """
    ["e4e2bcc3-3489-464c-bf06-c287c50b105e"]
    """
    if response.status_code != requests.codes.ok:
      self._handle_error_response(response)
      self._logger.error('Cancel order failed  %s, %s' %
                         (order.order_id, order.internal.exchange_order_id))
      try:
        message = response.json()
        if message.get('message') == "order not found":
          self._logger.info('order is dead already! %s', order.order_id)
          if not self._order_manager.is_dead_order(order):
            self._order_manager.set_order(order, cancel_confirmed=True)
      except ValueError:
        order.cancel_sent = False
      return OrderEvent.CANCEL_ERROR

    message = response.json()
    if 'message' in message:
      self._logger.error('Cancel order (%s) error! %s', order.order_id, message)
      if self._order_manager.is_dead_order(order):
        self._logger.warning('Order (%s) is dead already!', order.order_id)
      else:
        self._logger.error('Unexpected error! Need inspection! %s', message)
      order.cancel_sent = False
      return OrderEvent.CANCEL_ERROR
    elif message[0] != order.internal.exchange_order_id:
      self._logger.error('Inconsistent exchange order id %s\n%s', message, order)
      order.cancel_sent = False
      return OrderEvent.CANCEL_ERROR
    else:
      self._order_manager.set_order(order, cancel_confirmed=True)
      return OrderEvent.CANCEL_CONFIRMED

  def _update_order(self, order_info, is_rest):
    exchange_order_id = order_info.internal.exchange_order_id
    status = order_info.internal.status
    order = self._order_manager.get(exchange_order_id=exchange_order_id)
    if order is None:
      order_id = self._order_id_generator.gen_order_id(is_foreign=True)
      order_info.order_id = order_id
      self._order_manager.add_order(order_info)
      order = order_info
      if is_rest is False:
        self._logger.debug('WS message comes first!')

    if order.accepted is False:
      self._order_manager.set_order(order, exchange_order_id=exchange_order_id)

    if status in (GdaxOrderStatus.RECEIVED, GdaxOrderStatus.OPEN, GdaxOrderStatus.PENDING):
      pass
    elif status == GdaxOrderStatus.REJECTED:
      # treat rejected as cancel-confirmed
      if not self._order_manager.is_dead_order(order):
        self._order_manager.set_order(order, cancel_confirmed=True)
    elif status == GdaxOrderStatus.DONE_CANCELED:
      if not self._order_manager.is_dead_order(order):
        self._order_manager.set_order(order, fully_filled=True)
    elif status == GdaxOrderStatus.DONE_FILLED:
      if not self._order_manager.is_dead_order(order):
        self._order_manager.set_order(order, cancel_confirmed=True)
    else:
      raise ValueError('Unknown status %s', status)

  def _update_fill_info(self, order, price_last_fill, qty_last_fill, is_rest):
    if is_rest:
      fill_info = self._rest_fills[order.internal.exchange_order_id]
    else:
      fill_info = self._ws_fills[order.internal.exchange_order_id]

    if fill_info.value.qty_orig:
      is_same_qty = math.isclose(fill_info.value.qty_orig, order.qty)
      is_same_price = math.isclose(fill_info.value.price_orig, order.price)
      if not (is_same_qty and is_same_price):
        self._logger.error('Inconsistent fill_info. %s, %s', fill_info, order)
    else:
      fill_info.value.qty_orig = order.qty
      fill_info.value.price_orig = order.price

    last_fill = fill_info.update_by_last_fill(price_last_fill, qty_last_fill)
    if last_fill is None:
      self._logger.error('Invalid fill msg received! %s', order)
    return fill_info

  @exchange_message_handler(_MsgType.REST_ORDERS)
  def _handle_all_orders(self, exchange_message: PrivateExchangeMessage):
    response = exchange_message.data
    timestamp = exchange_message.received_timestamp
    if response.status_code != requests.codes.ok:
      self._handle_error_response(response)
      return

    message = response.json()
    for entry in message:
      native_symbol = entry['product_id']
      product = GdaxProduct.FromStrNativeProduct(native_symbol)
      if product not in self.og_info.products:
        continue  # ignore

      if entry['type'] == 'market':
        continue  # ignore

      # noinspection PyBroadException
      try:
        order_info = create_order_from_open_order_message(entry, timestamp)
        self._update_order(order_info, is_rest=True)
      except (KeyError, ValueError):
        self._logger.error('Parsing order message error: %s', json.dumps(entry, indent=2))
      except Exception:
        self._logger.exception('Parsing order message error: %s', json.dumps(entry, indent=2))

    if not self.is_ready():
      self._og_ready.set_orders_ready(True)
      self._og_ready.check_ready()

  @exchange_message_handler(_MsgType.REST_ACCOUNT_BALANCE)
  def _handle_rest_balance(self, exchange_message: PrivateExchangeMessage):
    response = exchange_message.data
    if response.status_code != requests.codes.ok:
      self._handle_error_response(response)
      return

    message = response.json()
    self._og_info.update_balance(message)
    self._og_logger.write_balance(self._og_info.balance_as_proto())
    if not self.is_ready():
      self._og_ready.set_balance_ready(True)
      self._og_ready.check_ready()

  @exchange_message_handler(_MsgType.REST_FILLS)
  def _handle_rest_fills(self, exchange_message: PrivateExchangeMessage):
    response = exchange_message.data
    timestamp = exchange_message.received_timestamp
    if response.status_code != requests.codes.ok:
      self._handle_error_response(response)
      return

    fill_list = response.json()
    is_first_fill_query = not bool(self._processed_fills)
    if is_first_fill_query:
      # Make it none empty, so after this we know it is not the first query.
      self._logger.info('First query!')
      self._processed_fills.add(None)

    for fill_msg in fill_list:
      trade_id = fill_msg['trade_id']
      exchange_order_id = str(fill_msg['order_id'])
      key = (trade_id, exchange_order_id)
      price_last_fill = float(fill_msg['price'])
      qty_last_fill = float(fill_msg['size'])
      fee = float(fill_msg['fee'])
      fill_type = None
      try:
        fill_type = getattr(coin_order, GdaxFillType(fill_msg['liquidity']).name)
      except (ValueError, AttributeError):
        self._logger.error('Unknown liquidity type: %s' % fill_msg)

      if is_first_fill_query:
        self._logger.info('old fill: %s', key)
        self._processed_fills.add(key)
        continue

      if key not in self._processed_fills:
        order = self._order_manager.get(exchange_order_id=exchange_order_id)
        if order is None:
          continue
        self._processed_fills.add(key)
        fill_info = self._update_fill_info(order, price_last_fill, qty_last_fill, is_rest=True)
        last_fill = order.internal.fill.update_by_accum_fill(fill_info.value.price_avg_fill,
                                                             fill_info.value.qty_accum_fill)
        order.internal.fee = fee + (order.internal.fee or 0.)
        if last_fill is not None:
          self._publish_fill(last_fill.price_last_fill, last_fill.qty_last_fill, order)
          self._og_logger.gen_order_filled_event_and_log(order, timestamp, fill_type)
          order.internal.fill.log_last_fill(order, self._logger)
          self._logger.warning('Rest fill comes first! %s, %s',
                               order.order_id,
                               order.internal.exchange_order_id)

  @exchange_message_handler(_MsgType.WS_HEART_BEAT)
  def _handle_ws_heart_beat(self, exchange_message: PrivateExchangeMessage):
    data = exchange_message.data
    self._logger.debug('%s', data)

  @exchange_message_handler(_MsgType.WS_USER)
  def _handle_ws_user(self, exchange_message: PrivateExchangeMessage):
    data = exchange_message.data
    timestamp = exchange_message.received_timestamp
    ws_msg_type = MessageType(data['type'])
    if ws_msg_type == MessageType.MATCH:
      exchange_order_id = data['maker_order_id']
      order = self._order_manager.get(exchange_order_id=exchange_order_id)
      fill_type = coin_order.MAKER_FILL_TYPE
      if order is None:
        exchange_order_id = data['taker_order_id']
        order = self._order_manager.get(exchange_order_id=exchange_order_id)
        fill_type = coin_order.TAKER_FILL_TYPE
        if order is None:
          self._logger.error('Order not found for this fill message %s', data)
          return
      qty_last_fill = float(data['size'])
      price_last_fill = float(data['price'])
      fill_info = self._update_fill_info(order, price_last_fill, qty_last_fill, is_rest=False)
      last_fill = order.internal.fill.update_by_accum_fill(fill_info.value.price_avg_fill,
                                                           fill_info.value.qty_accum_fill)
      if last_fill is not None:
        self._publish_fill(last_fill.price_last_fill, last_fill.qty_last_fill, order)
        self._og_logger.gen_order_filled_event_and_log(order, timestamp, fill_type=fill_type)
        order.internal.fill.log_last_fill(order, self._logger)
    elif ws_msg_type == MessageType.ACTIVATE:
      self._logger.error('Activate message received! %s', data)
    else:
      try:
        order_info = create_order_from_ws_order_update(data, exchange_message.received_timestamp)
        self._update_order(order_info, is_rest=False)
      except (KeyError, ValueError):
        self._logger.exception('Parsing ws order message error: %s', json.dumps(data, indent=2))
